"""
@author: chenzhenhua
@project: jf_fashion
@file: mnist.py
@time: 2021/8/2 0002 10:53
@desc:
"""

import torch
import torch.nn.functional as F
import torchvision
from torch.utils.data import DataLoader


class Mnist(object):
    """mnist数据集

    读取mnist数据，以及对mnist数据的训练和测试

    参数：
        data_path: 数据存放路径
            mnist数据存放的目录
        train_losses: 训练集的损失值

        train_counter：

        test_losses：测试集的损失值

    """

    def __init__(self, data_path, train_losses=[], train_counter=[], test_losses=[]):
        self.data_path = data_path
        self.train_losses = train_losses
        self.train_counter = train_counter
        self.test_losses = test_losses

    def train_data(self, batch_size_train):
        """

        :param batch_size_train: 训练集数据的批处理大小
        :return: 训练集数据
        """
        train_loader = torch.utils.data.DataLoader(
            torchvision.datasets.MNIST(self.data_path, train=True, download=True,
                                       transform=torchvision.transforms.Compose([
                                           torchvision.transforms.ToTensor(),
                                           torchvision.transforms.Normalize(
                                               (0.1307,), (0.3081,))
                                       ])),
            batch_size=batch_size_train, shuffle=True)
        return train_loader

    def test_data(self, batch_size_test):
        """

        :param batch_size_test: 测试集数据的批处理大小
        :return: 测试集数据
        """
        test_loader = torch.utils.data.DataLoader(
            torchvision.datasets.MNIST(self.data_path, train=False, download=True,
                                       transform=torchvision.transforms.Compose([
                                           torchvision.transforms.ToTensor(),
                                           torchvision.transforms.Normalize(
                                               (0.1307,), (0.3081,))
                                       ])),
            batch_size=batch_size_test, shuffle=True)
        return test_loader

    def train(self, network, optimizer, epoch, log_interval, batch_size_train, network_path='data/model.pth',
              optimizer_path='data/optimizer.pth'):
        """

        :param network: pytorch模型
        :param optimizer: 优化器
        :param epoch: 迭代次数
        :param log_interval:
        :param batch_size_train: 批处理大小
        :param network_path: 网络模型保存路径
        :param optimizer_path: 优化参数保存路径
        :return:
        """
        train_loader = self.train_data(batch_size_train)
        network.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = network(data)
            loss = F.nll_loss(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % log_interval == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                           100. * batch_idx / len(train_loader), loss.item()))
                self.train_losses.append(loss.item())
                self.train_counter.append(
                    (batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
                torch.save(network.state_dict(), network_path)
                torch.save(optimizer.state_dict(), optimizer_path)

    def test(self, network, batch_size_test):
        """

        :param network: pytorch模型
        :param batch_size_test: 批处理大小
        :return:
        """
        test_loader = self.test_data(batch_size_test)
        network.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                output = network(data)
                test_loss += F.nll_loss(output, target, size_average=False).item()
                pred = output.data.max(1, keepdim=True)[1]
                correct += pred.eq(target.data.view_as(pred)).sum()
        test_loss /= len(test_loader.dataset)
        self.test_losses.append(test_loss)
        print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))
